#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <string.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <arpa/inet.h>

#include "include.h"
#include "HelperUtils.h"

#include "ClientHandler.h"

void handleListMusic(int clientCommandSocket, char *documentRoot);
void handleListPictures(int clientCommandSocket, char *documentRoot);
void handleGet(int clientCommandSocket, const struct sockaddr_in clientAddress, char *documentRoot, char *getCommand);
void handlePut(int clientCommandSocket, const struct sockaddr_in clientAddress, char *documentRoot, char *getCommand);

int handleClient(int clientCommandSocket, const struct sockaddr_in clientAddress, char* documentRoot) {
    char receiveBuffer[RCVBUFSIZE];    // Buffer to hold received command
    int receivedBufferSize;
    
    while (1) {
        // Clear receivedBuffer
        memset(&receiveBuffer, 0, sizeof(char)*RCVBUFSIZE);

        // Assemble the full message string sent by client
        // Receive first message from client
        receivedBufferSize = recv(clientCommandSocket, receiveBuffer, RCVBUFSIZE-1, 0);
        if (receivedBufferSize < 0) {
            dieWithError("FATAL: recv() failed");
        } else if (receivedBufferSize == 0) {
            // Socket has been closed at client end
            // Server should close the connection
            #if DEBUG == 1
            fprintf(stdout, "DEBUG: Socket is closed on client side.\n");
            #endif
            break;
        }
        receiveBuffer[receivedBufferSize] = '\0';     // NULL pad the string
        
        int receiveBufferLength = strlen(receiveBuffer);
        // Finish receive when EOT character (0x4) is detected
        #if NC_MODE == 1
        while (receiveBuffer[strlen(receiveBuffer)-1] != '\n' ) {
        #else
        while (receiveBuffer[strlen(receiveBuffer)-1] != 0x4 ) {
        #endif
            receiveBufferLength = strlen(receiveBuffer);

            // See if there is more data to receive
            receivedBufferSize = recv(clientCommandSocket, receiveBuffer+receivedBufferSize, (RCVBUFSIZE-1)-receiveBufferLength, 0);
            if (receivedBufferSize < 0) {
                dieWithError("FATAL: recv() failed");
            } else if (receivedBufferSize == 0) {
                // Socket has been closed at client end
                // Server should close the connection
                #if DEBUG == 1
                fprintf(stdout, "DEBUG: Socket is closed on client side.\n");
                #endif
                break;
            }
            receiveBuffer[receivedBufferSize+receiveBufferLength] = '\0';
        }
        
        // Strip the EOT character
        receiveBuffer[strlen(receiveBuffer)-1] = '\0';

        // Find out which command is sent by client
        if ( strcmp(receiveBuffer, "LIST-MUSIC")==0 ) {
            #if DEBUG == 1
            fprintf(stderr, "DEBUG: LIST-MUSIC\n");
            #endif
            handleListMusic(clientCommandSocket, documentRoot);
        } else if ( strcmp(receiveBuffer, "LIST-PICTURES")==0 ) {
            #if DEBUG == 1
            fprintf(stderr, "DEBUG: LIST-PICTURES\n");
            #endif
            handleListPictures(clientCommandSocket, documentRoot);
        } else if ( strncmp(receiveBuffer, "GET", 3)==0 ) {
            #if DEBUG == 1
            fprintf(stderr, "DEBUG: GET\n");
            #endif
            handleGet(clientCommandSocket, clientAddress, documentRoot, receiveBuffer);
        } else if ( strncmp(receiveBuffer, "PUT", 3)==0 ) {
            #if DEBUG == 1
            fprintf(stderr, "DEBUG: PUT\n");
            #endif
            handlePut(clientCommandSocket, clientAddress, documentRoot, receiveBuffer);
        } else if ( strcmp(receiveBuffer, "BYE")==0 ) {
            #if DEBUG == 1
            fprintf(stderr, "DEBUG: BYE\n");
            #endif
            break;
        } else {
            // Unrecognized command
            #if INFO == 1
            fprintf(stdout, "INFO: Unrecognized command from client.\n");
            #endif
        }
    }

    #if DEBUG == 1
    fprintf(stderr, "DEBUG: Closing client connection.\n");
    #endif
    close(clientCommandSocket);
    return 0;
}

void handleListMusic(int clientCommandSocket, char *documentRoot) {
    // Operating System command to be run
    char* osCommand;
    // Allocate enough memory for osCommand: 
    //  size = length of "ls -l" + length of documentRoot 
    //          + length of "/PICTURES, or /MUSIC"
    // 32 ought to be enough here
    osCommand = (char *)malloc((32+strlen(documentRoot))*sizeof(char));

    // Setup OS command to be run
    //  ls -l $documentRoot/MUSIC
    strcpy(osCommand, "ls -l ");
    strcat(osCommand, documentRoot);
    strcat(osCommand, "/MUSIC");
    // Run osCommand and store result in osCommandPipe
    FILE *osCommandPipe = popen(osCommand, "r");
    if (osCommandPipe == NULL)
        dieWithError("FATAL: popen() failed.");

    // Send result
    char sendBuffer[CMD_SEND_BUF_SIZE];
    while (fgets(sendBuffer, sizeof(sendBuffer), osCommandPipe) != NULL) {
        if (send(clientCommandSocket, sendBuffer, strlen(sendBuffer), 0) != strlen(sendBuffer)) 
            dieWithError("FATAL: send() failed.");
    }
    // Append an EOT character to indicate end of transmission
    char eot = 0x4;
    if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");

    // Clean up
    pclose(osCommandPipe);
    free(osCommand);
}

void handleListPictures(int clientCommandSocket, char *documentRoot) {
    // Operating System command to be run
    char* osCommand;
    // Allocate enough memory for osCommand: 
    //  size = length of "ls -l" + length of documentRoot 
    //          + length of "/PICTURES, or /MUSIC"
    // 32 ought to be enough here
    osCommand = (char *)malloc((32+strlen(documentRoot))*sizeof(char));

    // Setup OS command to be run
    //  ls -l $documentRoot/MUSIC
    strcpy(osCommand, "ls -l ");
    strcat(osCommand, documentRoot);
    strcat(osCommand, "/PICTURES");
    // Run osCommand and store result in osCommandPipe
    FILE *osCommandPipe = popen(osCommand, "r");
    if (osCommandPipe == NULL)
        dieWithError("FATAL: popen() failed.");

    // Send result
    char sendBuffer[CMD_SEND_BUF_SIZE];
    while (fgets(sendBuffer, sizeof(sendBuffer), osCommandPipe) != NULL) {
        if (send(clientCommandSocket, sendBuffer, strlen(sendBuffer), 0) != strlen(sendBuffer)) 
            dieWithError("FATAL: send() failed.");
    }
    // Append an EOT character to indicate end of transmission
    char eot = 0x4;
    if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
        dieWithError("FATAL: send() failed.");
    
    // Clean up
    pclose(osCommandPipe);
    free(osCommand);
}

void handleGet(int clientCommandSocket, const struct sockaddr_in clientAddress, char *documentRoot, char *getCommand) {
    char eot = 0x4;
    char buffer[FILE_SEND_BUF_SIZE];

    // find out if GET command received is a valid one
    // format of GET command is:
    //  GET [PORT] [Relative path to file]
    char *token1, *token2;
    token1 = strchr(getCommand, ' ');      // Tokenize port number 
    // token1 is NULL means no space is found in getCommand
    // token1 is less than 2 characters means subsequent part is malformatted
    if (token1 == NULL || strlen(token1) < 2) {
        // Let client know it's a malformatted command
        strcpy(buffer, "Error: Malformatted GET command.\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        char eot = 0x4;
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Get out
        return;
    } else {
        token2 = strchr(token1+1, ' ');
        if (token2 == NULL || strlen(token2) < 2) {
            // Let client know it's a malformatted command
            strcpy(buffer, "Error: Malformatted GET command.\n");
            if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
                dieWithError("FATAL: send() failed.");
            // Append an EOT character to indicate end of transmission
            if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
                dieWithError("FATAL: send() failed.");
            // Get out
            return;
        }
    }

    // Obtain file name from token
    char *fileName = token2+1;
    // Obtain port number client is listening on from token
    char portString[16];
    memset(&portString, 0, sizeof(portString));
    strncpy(portString, token1+1, strlen(token1)-strlen(token2)-1);
    int clientDataPort = atoi(portString);
    // If atoi returns 0, something is wrong with portString sent by client
    if (clientDataPort == 0) {
        // Let client know it's a malformatted command
        strcpy(buffer, "Error: Malformatted GET command.\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Get out
        return;
    }

    // Full path to the file requested
    char *filePath;
    // Allocate enough memory for filePath: 
    //  size = length of documentRoot + filename
    filePath = (char *)malloc((strlen(documentRoot)+strlen(fileName)+32)*sizeof(char));
    // Setup filePath
    strcpy(filePath, documentRoot);
    strcat(filePath, "/");
    strcat(filePath, fileName);

    FILE *file = fopen(filePath, "r");
    
    // If client is requesting a non-existant file, let it know
    if (file == NULL) {
        strcpy(buffer, "Error: File ");
        strcat(buffer, fileName);
        strcat(buffer, " does not exist.\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        return;
    } else {
        // Send an EOT character to tell client start transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Everything looks good, try to connect to socket opened by client and send the file
        // Create data socket
        int clientDataSocket;
        if ((clientDataSocket = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
            dieWithError("FATAL: socket() failed.");
        // Create client address
        struct sockaddr_in clientDataAddress;
        memset(&clientDataAddress, 0, sizeof(clientDataAddress));
        clientDataAddress.sin_family = AF_INET;
        clientDataAddress.sin_addr.s_addr = clientAddress.sin_addr.s_addr;
        clientDataAddress.sin_port = htons(clientDataPort);
        // Establish the connection
        if (connect(clientDataSocket, (struct sockaddr *)&clientDataAddress, sizeof(clientDataAddress)) < 0) {
            // Let client know attempt to connect to port provided has failed
            strcpy(buffer, "Error: Failed to connect to port ");
            strcat(buffer, portString);
            strcat(buffer, ".\n");
            if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
                dieWithError("FATAL: send() failed.");
            // Append an EOT character to indicate end of transmission
            if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
                dieWithError("FATAL: send() failed.");
            // Get out
            return;
        }

        // Send file in FILE_SEND_BUF_SIZE chunk
        int numberRead = 0;
        while ((numberRead = fread(buffer, 1, FILE_SEND_BUF_SIZE, file)) == FILE_SEND_BUF_SIZE) {
            if (send(clientDataSocket, buffer, numberRead, 0) != numberRead) 
                dieWithError("FATAL: send() failed.");
        }
        if (send(clientDataSocket, buffer, numberRead, 0) != numberRead) 
            dieWithError("FATAL: send() failed.");

        fclose(file);
        close(clientDataSocket);
    }
    
    // Clean up
    free(filePath);
}

void handlePut(int clientCommandSocket, const struct sockaddr_in clientAddress, char *documentRoot, char *getCommand) {
    char eot = 0x4;
    char buffer[FILE_SEND_BUF_SIZE];

    // find out if PUT command received is a valid one
    // format of PUT command is:
    //  PUT [PORT] [Relative path to file]
    char *token1, *token2;
    token1 = strchr(getCommand, ' ');      // Tokenize port number 
    // token1 is NULL means no space is found in getCommand
    // token1 is less than 2 characters means subsequent part is malformatted
    if (token1 == NULL || strlen(token1) < 2) {
        // Let client know it's a malformatted command
        strcpy(buffer, "Error: Malformatted GET command.\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Get out
        return;
    } else {
        token2 = strchr(token1+1, ' ');
        if (token2 == NULL || strlen(token2) < 2) {
            // Let client know it's a malformatted command
            strcpy(buffer, "Error: Malformatted GET command.\n");
            if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
                dieWithError("FATAL: send() failed.");
            // Append an EOT character to indicate end of transmission
            if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
                dieWithError("FATAL: send() failed.");
            // Get out
            return;
        }
    }
    
    // Obtain file name from token
    char *fileName = token2+1;
    // Obtain port number client is listening on from token
    char portString[16];
    memset(&portString, 0, sizeof(portString));
    strncpy(portString, token1+1, strlen(token1)-strlen(token2)-1);
    int clientDataPort = atoi(portString);
    // If atoi returns 0, something is wrong with portString sent by client
    if (clientDataPort == 0) {
        // Let client know it's a malformatted command
        strcpy(buffer, "Error: Malformatted GET command.\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Get out
        return;
    }

    // Full path to the file requested
    char *filePath;
    // Allocate enough memory for filePath: 
    //  size = length of documentRoot + filename
    filePath = (char *)malloc((strlen(documentRoot)+strlen(fileName)+32)*sizeof(char));
    // Setup filePath
    strcpy(filePath, documentRoot);
    strcat(filePath, "/");
    strcat(filePath, fileName);

    // Open file for writing
    // TODO: If client request to put a file with same name, it will be overwritten. Should list this in assumption.
    FILE *file = fopen(filePath, "w");
    if (file == NULL) {
        strcpy(buffer, "Error: Cannot create ");
        strcat(buffer, fileName);
        strcat(buffer, ".\n");
        if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
            dieWithError("FATAL: send() failed.");
        // Append an EOT character to indicate end of transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        return;
    } else {
        // Send an EOT character to tell client to start transmission
        if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
            dieWithError("FATAL: send() failed.");
        // Everything looks good, try to connect to socket opened by client and receive the file
        // Create data socket
        int clientDataSocket;
        if ((clientDataSocket = socket(PF_INET, SOCK_STREAM, IPPROTO_TCP)) < 0)
            dieWithError("FATAL: socket() failed.");
        // Create client address
        struct sockaddr_in clientDataAddress;
        memset(&clientDataAddress, 0, sizeof(clientDataAddress));
        clientDataAddress.sin_family = AF_INET;
        clientDataAddress.sin_addr.s_addr = clientAddress.sin_addr.s_addr;
        clientDataAddress.sin_port = htons(clientDataPort);
        // Establish the connection
        if (connect(clientDataSocket, (struct sockaddr *)&clientDataAddress, sizeof(clientDataAddress)) < 0) {
            // Let client know attempt to connect to port provided has failed
            strcpy(buffer, "Error: Failed to connect to port ");
            strcat(buffer, portString);
            strcat(buffer, ".\n");
            if (send(clientCommandSocket, buffer, strlen(buffer), 0) != strlen(buffer)) 
                dieWithError("FATAL: send() failed.");
            // Append an EOT character to indicate end of transmission
            if (send(clientCommandSocket, &eot, sizeof(eot), 0) != sizeof(eot))
                dieWithError("FATAL: send() failed.");
            // Get out
            return;
        }

        // Receive file in FILE_SEND_BUF_SIZE chunk
        // After sending the whole file, client is expected to close connection.
        int byteReceived;
        int numberWritten = 0;
        while ((byteReceived = recv(clientDataSocket, buffer, FILE_SEND_BUF_SIZE, 0)) > 0) {
            // Write the buffer to file
            if ((numberWritten = fwrite(buffer, 1, byteReceived, file)) != byteReceived)
                dieWithError("FATAL: fwrite() failed.");
        }
        fclose(file);
        close(clientDataSocket);
    }
    
    // Clean up
    free(filePath);
}

